import pandas as pd
import numpy as np
from pathlib import Path
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler, OneHotEncoder, OrdinalEncoder
from sklearn.utils.multiclass import type_of_target


class DataProcessor:
    def __init__(self, config, logger):
        self.config = config
        self.logger = logger
        self.preprocessor = None
        self.numeric_features = None
        self.categorical_features = None
        self.binning_features = None  # 分箱字段
        self.binning_config = config.get("binning", {})

    def load_data(self) -> tuple:
        """加载训练集和测试集，返回特征和标签"""
        data_dir = Path(self.config.get("paths.data_dir"))
        train_path = data_dir.joinpath(self.config.get("paths.train_file"))
        test_path = data_dir.joinpath(self.config.get("paths.test_file"))
        target_col = self.config.get("global_params.target_column")

        try:
            self.logger.info(f"加载训练集: {train_path}")
            train_df = pd.read_csv(train_path)
            self.logger.info(f"加载测试集: {test_path}")
            test_df = pd.read_csv(test_path)

            # 检查缺失值
            self._log_missing_values(train_df, "训练集")
            self._log_missing_values(test_df, "测试集")

            # 分离特征与标签
            X_train = train_df.drop(target_col, axis=1)
            y_train = train_df[target_col]
            X_test = test_df.drop(target_col, axis=1)
            y_test = test_df[target_col]

            self.logger.info(f"数据加载完成 - 训练集: {len(X_train)}样本, 测试集: {len(X_test)}样本")
            return X_train, y_train, X_test, y_test

        except FileNotFoundError as e:
            self.logger.error(f"数据文件不存在: {e.filename}")
            raise

    def _log_missing_values(self, df: pd.DataFrame, name: str):
        """记录缺失值信息"""
        missing = df.isnull().sum()[df.isnull().sum() > 0]
        if not missing.empty:
            self.logger.info(f"{name}缺失值情况:\n{missing}")
        else:
            self.logger.info(f"{name}无缺失值")

    def _create_binning_transformer(self, X_train: pd.DataFrame):
        """创建分箱转换器"""
        if not self.binning_config.get("enabled", False):
            return None, []

        binning_fields = self.binning_config.get("fields", {})
        if not binning_fields:
            return None, []

        # 验证分箱字段是否存在于训练数据中
        valid_fields = []
        for field in binning_fields:
            if field in X_train.columns:
                valid_fields.append(field)
            else:
                self.logger.warning(f"分箱配置中的字段'{field}'不存在于训练数据中，将被忽略")

        if not valid_fields:
            return None, []

        self.binning_features = valid_fields
        self.logger.info(f"分箱字段({len(self.binning_features)}个): {self.binning_features}")

        # 创建分箱管道
        transformers = []
        for field in valid_fields:
            cfg = binning_fields[field]
            field_pipeline = Pipeline(steps=[
                ('imputer', SimpleImputer(strategy='median')),  # 先填充缺失值
                ('binner', BinningTransformer(cfg)),  # 分箱转换
                ('encoder', OrdinalEncoder(handle_unknown='use_encoded_value', unknown_value=-1))  # 分箱后编码
            ])
            transformers.append((f'bin_{field}', field_pipeline, [field]))

        return ColumnTransformer(transformers=transformers), valid_fields

    def build_preprocessor(self, X_train: pd.DataFrame):
        """构建预处理管道，包含分箱处理"""
        # 处理分箱
        binning_transformer, binning_fields = self._create_binning_transformer(X_train)

        # 区分数值和分类特征（排除分箱字段）
        remaining_features = [col for col in X_train.columns if col not in binning_fields]
        self.numeric_features = X_train[remaining_features].select_dtypes(include=['int64', 'float64']).columns.tolist()
        self.categorical_features = X_train[remaining_features].select_dtypes(
            include=['object', 'category']).columns.tolist()

        self.logger.info(f"数值特征({len(self.numeric_features)}个): {self.numeric_features}")
        self.logger.info(f"分类特征({len(self.categorical_features)}个): {self.categorical_features}")

        # 数值特征管道
        numeric_transformer = Pipeline(steps=[
            ('imputer', SimpleImputer(strategy='median')),
            ('scaler', StandardScaler())
        ])

        # 分类特征管道
        categorical_transformer = Pipeline(steps=[
            ('imputer', SimpleImputer(strategy='most_frequent')),
            ('encoder', OneHotEncoder(drop='first', sparse_output=False, handle_unknown='ignore'))
        ])

        # 构建所有转换器列表
        transformers = []
        if binning_transformer:
            transformers.append(('binning', binning_transformer, binning_fields))
        transformers.extend([
            ('num', numeric_transformer, self.numeric_features),
            ('cat', categorical_transformer, self.categorical_features)
        ])

        # 组合预处理
        self.preprocessor = ColumnTransformer(transformers=transformers)
        return self.preprocessor

    @staticmethod
    def infer_task_type(y: pd.Series) -> str:
        """推断任务类型：binary/multiclass/regression"""
        target_type = type_of_target(y)
        if target_type == "continuous":
            return "regression"
        elif target_type == "binary":
            return "binary"
        elif target_type in ["multiclass", "multiclass-multioutput"]:
            return "multiclass"
        else:
            raise ValueError(f"不支持的标签类型: {target_type}")


class BinningTransformer:
    """自定义分箱转换器，支持多种分箱方式"""

    def __init__(self, config):
        self.config = config
        self.bins_ = None  # 存储拟合后的分箱边界
        self.labels_ = config.get("labels")  # 分箱标签

    def fit(self, X, y=None):
        """根据配置拟合分箱器"""
        # X是二维数组，需要转换为一维
        X_flat = X.flatten()
        bin_type = self.config.get("type", "custom")

        if bin_type == "custom" and "bins" in self.config:
            self.bins_ = self.config["bins"]

        elif bin_type == "qcut":
            q = self.config.get("q", 5)
            _, self.bins_ = pd.qcut(
                X_flat[~pd.isna(X_flat)],  # 排除缺失值
                q=q,
                labels=self.labels_,
                retbins=True,
                duplicates='drop'
            )

        elif bin_type == "cut":
            bins = self.config.get("bins", 5)
            _, self.bins_ = pd.cut(
                X_flat[~pd.isna(X_flat)],
                bins=bins,
                labels=self.labels_,
                retbins=True,
                duplicates='drop'
            )
        return self

    def transform(self, X):
        """应用分箱转换"""
        X_flat = X.flatten()
        if self.bins_ is None:
            raise ValueError("分箱器尚未拟合，请先调用fit方法")

        # 应用分箱
        binned = pd.cut(
            X_flat,
            bins=self.bins_,
            labels=self.labels_,
            include_lowest=True
        )

        # 转换为数值编码
        return binned.codes.reshape(-1, 1)